#!/usr/bin/env python3
# coding: utf-8

from __future__ import print_function, division, unicode_literals
from bst import Node
import bst


__all__ = ['RBTree', 'left_rotate', 'right_rotate']

RED = 'red'
BLACK = 'black'


class RBTree(bst.Tree):
    """
    Red-Black Tree
    """

    def insert(self, node, parent=None):
        """
        (see bst.Tree.insert.__doc__)
        """
        node = bst.Tree.insert(self, node, parent)
        if self.root is node:
            node.color = BLACK
        else:
            node.color = RED
            #self.insert_case2(node)  # tree root was just checked, skip case1
            while node is not self.root and node.p.color != BLACK:
                uncle = get_uncle(node)
                if uncle is not None and uncle.color == RED:
                    node.p.color = BLACK
                    uncle.color = BLACK
                    gp = get_grandparent(node)
                    if gp is not None:
                        gp.color = RED
                    node = gp
                    # next iteration of the main loop
                else:
                    gp = get_grandparent(node)
                    if node is node.p.right and node.p is gp.left:
                        left_rotate(node.p)
                        node = node.left
                    elif node is node.p.left and node.p is gp.right:
                        right_rotate(node.p)
                        node = node.right
                    gp = get_grandparent(node)
                    node.p.color = BLACK
                    if gp is not None:
                        gp.color = RED
                    if node is node.p.left:
                        right_rotate(gp)
                    else:
                        left_rotate(gp)
                    if self.root is gp:  # XXX
                        self.root = self.root.p
                    break
        return node

    def insert_case1(self, node):
        if node.p is None:
            node.color = BLACK
        else:
            if node.p.color != BLACK:
                uncle = get_uncle(node)
                if uncle is not None and uncle.color == RED:
                    node.p.color = BLACK
                    uncle.color = BLACK
                    grandp = get_grandparent(node)
                    grandp.color = RED
                    self.insert_case1(grandp)  # TODO: rewrite as loop
                else:
                    grandp = get_grandparent(node)
                    if node is node.p.right and node.p is grandp.left:
                        left_rotate(node.p)
                        node = node.left
                    elif node is node.p.left and node.p is grandp.right:
                        right_rotate(node.p)
                        node = node.right
                    grandp = get_grandparent(node)
                    node.p.color = BLACK
                    grandp.color = RED
                    if node is node.p.left:
                        right_rotate(grandp)
                    else:
                        left_rotate(grandp)
                    if self.root is grandp:  # XXX
                        self.root = self.root.p



def left_rotate(root):
    node = root.right
    root.right = node.left
    if node.left:
        node.left.p = root
    node.p = root.p
    if root.p and root.p.left is root:
        root.p.left = node
    elif root.p:
        root.p.right = node
    root.p = node
    node.left = root


def right_rotate(root):
    node = root.left
    root.left = node.right
    if node.right:
        node.right.p = root
    node.p = root.p
    if root.p and root.p.left is root:
        root.p.left = node
    elif root.p:
        root.p.right = node
    root.p = node
    node.right = root


def get_grandparent(node):
    if node and node.p:
        return node.p.p
    else:
        return None


def get_uncle(node):
    grandparent = get_grandparent(node)
    if grandparent is None:
        return None
    if node.p is grandparent.left:
        return grandparent.right
    else:
        return grandparent.left


def main():
    import numpy, sys, time
    times = {}
    for exp in range(1, 8):
        times[exp] = [None] * 2
        t = RBTree()
        l = list(range(1, 10 ** exp + 1))
        numpy.random.shuffle(l)

        ts = time.time()
        t = RBTree([Node(x) for x in l])
        times[exp][0] = time.time() - ts
        #print('{0:.5f} seconds to create a red-black tree from 10 ** {1} '
        #        'elements (unsorted)'.format(time.time() - ts, exp))

        ts = time.time()
        for x in l:
            assert t.find(x), 'tree corrupted'
        times[exp][1] = time.time() - ts
        #print('{0:.5f} seconds to locate 10 ** {1} elements'.format(time.time() - ts, exp))
    for exp in sorted(times.keys()):
        print('10 ** {0}: creation = {1:9.5f}, avg insert = {4:9.7f}, search * n = {2:9.5f}, '
                'search = {3:12.10f}'.format(exp, times[exp][0], times[exp][1],
                    times[exp][1] / 10 ** exp, times[exp][0] / 10 ** exp))
    print()
    for exp in sorted(times.keys())[1:]:
        print('creation, 10 ** {0} / 10 ** {1} = {2}'.format(exp, exp - 1,
            times[exp][0] / times[exp - 1][0]))
    for exp in sorted(times.keys())[1:]:
        print('search * n, 10 ** {0} / 10 ** {1} = {2}'.format(exp, exp - 1,
            times[exp][1] / times[exp - 1][1]))


if __name__ == "__main__":
    main()
